package com.xiashitech.agent.instrumentation.db.common;

import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import lombok.*;

import java.util.HashMap;
import java.util.Map;


public class XSUpdateSizeContext {

    @Getter
    @Setter
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class Request {
        String sql;
        int size;
    }

    /**
     * 保存execute方法上下文
     */
    @Getter
    @Setter
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class UpdateSizeContext {
        boolean fromAddBatch;
        boolean executeType;
        Request request;
        Context context;
        Scope scope;
    }

    /**
     * 通过Statement获取UpdateSizeContext
     * @param statement
     * @return
     */
    public static UpdateSizeContext getUpdateSizeContext(Object statement) {
        return statementWithContext.get().get(statement);
    }

    /**
     * 设置UpdateSizeContext，key为Statement
     * @param statement
     * @param updateSizeContext
     */
    public static void setStatementWithContext(Object statement,UpdateSizeContext updateSizeContext) {
        statementWithContext.get().put(statement,updateSizeContext);
    }

    private static ThreadLocal<Map<Object,UpdateSizeContext>> statementWithContext = new ThreadLocal<Map<Object,UpdateSizeContext>>() {
        protected Map<Object,UpdateSizeContext> initialValue() {
            return new HashMap<Object,UpdateSizeContext>();
        }
    };

    public static void refreshCurrentThreadLocal() {
        statementWithContext.set(new HashMap<Object,UpdateSizeContext>());
    }
}
